from sqlalchemy import create_engine, text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
import os
import pymysql

# 数据库连接URL配置
# 如果是生产环境，可以从环境变量读取
DB_USER = os.getenv("DB_USER", "root")
DB_PASSWORD = os.getenv("DB_PASSWORD", "123456")
# DB_PASSWORD = os.getenv("DB_PASSWORD", "rm-2ze8t37k3h4d54586")
DB_HOST = os.getenv("DB_HOST", "localhost")
DB_PORT = os.getenv("DB_PORT", "3306")
DB_NAME = os.getenv("DB_NAME", "merchant_collection")

# 创建数据库（如果不存在）
def create_database_if_not_exists():
    try:
        # 创建不带数据库名的连接
        conn = pymysql.connect(
            host=DB_HOST,
            user=DB_USER,
            password=DB_PASSWORD,
            port=int(DB_PORT)
        )
        cursor = conn.cursor()
        
        # 创建数据库（如果不存在）
        cursor.execute(f"CREATE DATABASE IF NOT EXISTS {DB_NAME} CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
        conn.commit()
        cursor.close()
        conn.close()
        print(f"数据库 {DB_NAME} 检查/创建成功")
    except Exception as e:
        print(f"创建数据库时出错: {str(e)}")
        raise

# 创建数据库
create_database_if_not_exists()

# 完整的数据库URL
SQLALCHEMY_DATABASE_URL = f"mysql+pymysql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}"

# 创建SQLAlchemy引擎
engine = create_engine(
    SQLALCHEMY_DATABASE_URL,
    echo=False,  # 关闭SQL语句输出
    pool_pre_ping=True,  # 启用连接池心跳检测
    pool_recycle=3600,  # 连接回收时间（秒）
)

# 创建数据库会话类
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# 创建基类
Base = declarative_base()

# 获取数据库会话的依赖
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

# 初始化数据库表
def init_db():
    """初始化数据库表"""
    from models import Base
    Base.metadata.create_all(bind=engine)
    print("数据库表初始化完成")

def create_merchant_data_table(table_name: str):
    """为商户创建数据表
    
    Args:
        table_name: 表名，格式为 商户名_商户ID
    """
    # 创建表的SQL语句
    create_table_sql = f"""
    CREATE TABLE IF NOT EXISTS {table_name} (
        id INT AUTO_INCREMENT PRIMARY KEY,
        data_period VARCHAR(50) NOT NULL COMMENT '数据周期',
        topic VARCHAR(100) NOT NULL COMMENT '主题',
        group_name VARCHAR(100) NOT NULL COMMENT '分组',
        metric_name VARCHAR(100) NOT NULL COMMENT '指标名',
        metric_subitem VARCHAR(100) NOT NULL COMMENT '指标子项',
        metric_value VARCHAR(200) COMMENT '指标值',
        remark TEXT COMMENT '备注',
        created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
        modified_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '修改时间',
        UNIQUE KEY uk_metric (topic, group_name, metric_name, metric_subitem, data_period)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
    """
    
    try:
        # 执行创建表
        with engine.connect() as conn:
            conn.execute(text(create_table_sql))
            conn.commit()
        print(f"表 {table_name} 创建成功")
    except Exception as e:
        print(f"创建表 {table_name} 时出错: {str(e)}")
        raise 